#include <sstream>
#include <iomanip>
#include <opencv2/dnn.hpp>

#include "base_util/utils.h"
#include "scrfd_predictor.h"

namespace ai {


static void dump_tensor_attr(rknn_tensor_attr *attr) {
  printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
          "zp=%d, scale=%f\n",
          attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
          attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
          get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}
static inline float fast_exp(float x) {
  // return exp(x);
  union
  {
    uint32_t i;
    float f;
  } v;
  v.i = (12102203.1616540672 * x + 1064807160.56887296);
  return v.f;
}
static float DeQnt2F32(int8_t qnt, int zp, float scale) {
  return ((float)qnt - (float)zp) * scale;
}

RKScrfdPredictor::RKScrfdPredictor(RKPackPredictor* model, RKModelManager* manager, LogInfo *lg):
  RKPredictor(model, manager, lg) { }


int RKScrfdPredictor::start_postprocess_thread_imp() {
  int ret = 0;
  std::pair<std::pair<int,int>,int> cur_block;
  // stop_flag2 = common::CameraGrabing;
  // while (stop_flag2 == common::CameraGrabing) {

  //   run_mtx2.lock();
  //   if (post_data2.empty()) {run_mtx2.unlock(); std::this_thread::sleep_for(std::chrono::milliseconds(3));continue;}
  //   cur_block = post_data2.front();
  //   post_data2.pop();
  //   run_mtx2.unlock();

  //   int loop_idx = cur_block.first.second;
  //   std::vector<BaseInfo*> det_infos;
  //   ret = det_hb_->PostProcess(det_infos, cur_block.first.first, this, cur_block.second);  //处理图片
  //   run_mtx3.lock();
  //   out_data3[loop_idx].push(det_infos);  //将处理后的数据返回到det_infos,加入队列中
  //   run_mtx3.unlock();
  // }
  // stop_flag2 = common::CameraOpened;
  return 0;
}

int RKScrfdPredictor::RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) {
  int ret;
  // printf("[liufeng debug] RKScrfdPredictor::RunDet:\n");
  stream::ImageBlob img_blob(stream::ImageBlobMode_BGR);
  /*
  if (!mdl_rk->cfg->transforms->run(blob->img, img_blob, mdl_rk->cfg)) { 
    printf("transforms->run fail \n");
    return model_image_channels_check_error;
  }
  */
  // std::cout << "img_blob.img w:" << img_blob.img.cols << " h:" << img_blob.img.rows << std::endl;
  // cv::imwrite("img_blob.img.jpg",img_blob.img);
  // cv::imwrite("blobimg.jpg",blob->img);

  img_blob.ori_im_shape = blob->ori_im_shape;

  cv::Mat padimg;
  SCRFDScaleParams scale_params;
  resize_unscale(blob->img, padimg, mdl_rk->cfg->input_shape[1], mdl_rk->cfg->input_shape[0], scale_params);// h w
  // cv::imwrite("padimg.jpg",padimg);
  // printf("[liufeng debug] scale_params:%4d %4d %.4f %4d\n",scale_params.dw,scale_params.dh,scale_params.ratio,scale_params.flag);

  rknn_tensor_attr output_attrs[mdl_rk->output_count];
  memset(output_attrs, 0, sizeof(output_attrs));
  for (int i = 0; i < mdl_rk->output_count; i++) {
    output_attrs[i].index = i;
    ret = rknn_query(mdl_rk->ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
    // dump_tensor_attr(&(output_attrs[i]));
  } 

  // mdl_rk->inputs[0].buf = (void *)(img_blob.img.data);
  mdl_rk->inputs[0].buf = (void *)(padimg.data);
  rknn_inputs_set(mdl_rk->ctx, mdl_rk->input_count, &mdl_rk->inputs[0]);

  rknn_core_mask core_mask = RKNN_NPU_CORE_0;
  if(mdl_rk->cfg->gpu_id==1) { core_mask = RKNN_NPU_CORE_1; }
  if(mdl_rk->cfg->gpu_id==2) { core_mask = RKNN_NPU_CORE_2; }
  ret = rknn_set_core_mask(mdl_rk->ctx, core_mask);
  ret = rknn_run(mdl_rk->ctx, NULL);
  if (ret != 0) {printf("rknn_run fail %d", ret); return -1;}
  ret = rknn_outputs_get(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0], NULL);
  if (ret != 0) {printf("rknn_outputs_get fail %d", ret); return -1;}
 
  // 后处理部分
  std::vector<float> out_scales;
  std::vector<int32_t> out_zps;
  for (int i = 0; i < mdl_rk->output_count; ++i) {
    out_scales.push_back(output_attrs[i].scale);
    out_zps.push_back(output_attrs[i].zp);
  }

  float* output_tensors[mdl_rk->output_count];
  for (int j = 0; j < mdl_rk->output_count; j++)
  {

      int8_t *pblob = (int8_t *)mdl_rk->outputs[j].buf;
      float* pf=(float*)malloc(output_attrs[j].n_elems*sizeof(float));
      
      for(int i=0;i<output_attrs[j].n_elems;i++)
      {
          pf[i]=DeQnt2F32(pblob[i], out_zps[j], out_scales[j]);               
      }
      output_tensors[j]=pf;

  }

  std::vector<BoxfWithLandmarks> detected_boxes;
  //std::vector<FaceDetInfo*> det_infos_tmp;  
  GetConvDetectionResult(output_tensors, scale_params, mdl_rk->cfg->draw_threshold,  mdl_rk->cfg->nms_threshold, blob->img.rows, blob->img.cols, det_infos, 400, detected_boxes);
  for(int i=0;i<mdl_rk->output_count;i++)
  {
      free(output_tensors[i]);
  }



  rknn_outputs_release(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0]);

  return 0;

}

void RKScrfdPredictor::resize_unscale(const cv::Mat &mat, cv::Mat &mat_rs,
                           int target_height, int target_width,
                           SCRFDScaleParams &scale_params)
{
  if (mat.empty()) return;
  int img_height = static_cast<int>(mat.rows);
  int img_width = static_cast<int>(mat.cols);

  mat_rs = cv::Mat(target_height, target_width, CV_8UC3,
                   cv::Scalar(0, 0, 0));
  // scale ratio (new / old) new_shape(h,w)
  float w_r = (float) target_width / (float) img_width;
  float h_r = (float) target_height / (float) img_height;
  float r = std::min(w_r, h_r);
  // compute padding
  int new_unpad_w = static_cast<int>((float) img_width * r); // floor
  int new_unpad_h = static_cast<int>((float) img_height * r); // floor
  int pad_w = target_width - new_unpad_w; // >=0
  int pad_h = target_height - new_unpad_h; // >=0

  int dw = pad_w / 2;
  int dh = pad_h / 2;

  // resize with unscaling
  cv::Mat new_unpad_mat;
  // cv::Mat new_unpad_mat = mat.clone(); // may not need clone.
  cv::resize(mat, new_unpad_mat, cv::Size(new_unpad_w, new_unpad_h));
  new_unpad_mat.copyTo(mat_rs(cv::Rect(dw, dh, new_unpad_w, new_unpad_h)));

  // record scale params.
  scale_params.ratio = r;
  scale_params.dw = dw;
  scale_params.dh = dh;
  scale_params.flag = true;
}

int RKScrfdPredictor::GetConvDetectionResult(float** output_tensors, SCRFDScaleParams scale_params, float score_threshold, float iou_threshold, int img_height, int img_width, std::vector<BaseInfo*>& det_infos, unsigned int topk, std::vector<BoxfWithLandmarks>& detected_boxes)
{

    std::vector<BoxfWithLandmarks> bbox_kps_collection;
    generate_bboxes_kps(scale_params, bbox_kps_collection, output_tensors,
                        score_threshold, img_height, img_width, mdl_rk->cfg->input_shape[1], mdl_rk->cfg->input_shape[0]);
    // printf("[liufeng debug+] bbox_kps_collection.size:%d\n",bbox_kps_collection.size());

    nms_bboxes_kps(bbox_kps_collection, detected_boxes, iou_threshold, topk);

    // printf("[liufeng debug] detected_boxes.size:%d\n",detected_boxes.size());


    std::vector<int> indices;
    for (int idx=0; idx < detected_boxes.size(); idx++) {indices.push_back(idx);}

    for (auto& idx : indices) {
      //box左上角点 //box右下角点
      DetInfo* di = new DetInfo(detected_boxes[idx].box.score, detected_boxes[idx].box.label, cv::Rect(cv::Point2f(detected_boxes[idx].box.x1,detected_boxes[idx].box.y1), cv::Point2f(detected_boxes[idx].box.x2,detected_boxes[idx].box.y2)));
      di->category = "face";
      di->points = detected_boxes[idx].landmarks.points;//5个关键点
      det_infos.emplace_back(di);
    }    


    return 0;
}

//SCRFD人脸检测、关键点检测后处理算法具体函数实现
void RKScrfdPredictor::generate_points(const int target_height, const int target_width)
{
  if (center_points_is_update) return;
  // 8, 16, 32
  for (auto stride : feat_stride_fpn)
  {
    unsigned int num_grid_w = target_width / stride;
    unsigned int num_grid_h = target_height / stride;
    // y
    for (unsigned int i = 0; i < num_grid_h; ++i)
    {
      // x
      for (unsigned int j = 0; j < num_grid_w; ++j)
      {
        // num_anchors, col major
        for (unsigned int k = 0; k < num_anchors; ++k)
        {
          SCRFDPoint point;
          point.cx = (float) j;
          point.cy = (float) i;
          point.stride = (float) stride;
          center_points[stride].push_back(point);
        }

      }
    }
  }

  center_points_is_update = true;
}

void RKScrfdPredictor::generate_bboxes_kps_single_stride(const SCRFDScaleParams &scale_params,
                                           float* score_ptr,
                                           float* bbox_ptr,
                                           float* kps_ptr,
                                           unsigned int stride,
                                           float score_threshold,
                                           float img_height,
                                           float img_width,
                                           std::vector<BoxfWithLandmarks> &bbox_kps_collection)
{

  unsigned int nms_pre_ = (stride / 8) * nms_pre; // 1 * 1000,2*1000,...
  nms_pre_ = nms_pre_ >= nms_pre ? nms_pre_ : nms_pre;
  // printf("[liufeng debug] nms_pre_: %d\n",nms_pre_);

  unsigned int num_points;
  if(stride == 8)
  {
    num_points = output_num_points[0];
  }
  else if (stride == 16)
  {
    num_points = output_num_points[1];
  }
  else
  {
    num_points = output_num_points[2];
  }
  // printf("[liufeng debug] num_points:%d score_threshold:%.4f\n",num_points,score_threshold);

  float ratio = scale_params.ratio;
  int dw = scale_params.dw;
  int dh = scale_params.dh;
  // printf("[liufeng debug] ratio:%.4f dw:%d dh:%d\n",ratio,dw,dh);
  unsigned int count = 0;
  auto &stride_points = center_points[stride];

  for (unsigned int i = 0; i < num_points; ++i)
  {
    float cls_conf = score_ptr[i];
    // printf("[liufeng debug] cls_conf:%.4f \n",cls_conf);
    if (cls_conf < score_threshold) continue; // filter
    auto &point = stride_points.at(i);
    const float cx = point.cx; // cx
    const float cy = point.cy; // cy
    const float s = point.stride; // stride

    // bbox
    const float *offsets = bbox_ptr + i * 4;
    float l = offsets[0]; // left
    float t = offsets[1]; // top
    float r = offsets[2]; // right
    float b = offsets[3]; // bottom

    BoxfWithLandmarks box_kps;
    float x1 = ((cx - l) * s - (float) dw) / ratio;  // cx - l x1
    float y1 = ((cy - t) * s - (float) dh) / ratio;  // cy - t y1
    float x2 = ((cx + r) * s - (float) dw) / ratio;  // cx + r x2
    float y2 = ((cy + b) * s - (float) dh) / ratio;  // cy + b y2
    box_kps.box.x1 = std::max(0.f, x1);
    box_kps.box.y1 = std::max(0.f, y1);
    box_kps.box.x2 = std::min(img_width - 1.f, x2);
    box_kps.box.y2 = std::min(img_height - 1.f, y2);
    box_kps.box.score = cls_conf;
    box_kps.box.label = 1;
    box_kps.box.label_text = "face";
    box_kps.box.flag = true;

    // landmarks
    const float *kps_offsets = kps_ptr + i * 10;
    for (unsigned int j = 0; j < 10; j += 2)
    {
      cv::Point2f kps;
      float kps_l = kps_offsets[j];
      float kps_t = kps_offsets[j + 1];
      float kps_x = ((cx + kps_l) * s - (float) dw) / ratio;  // cx + l x
      float kps_y = ((cy + kps_t) * s - (float) dh) / ratio;  // cy + t y
      kps.x = std::min(std::max(0.f, kps_x), img_width - 1.f);
      kps.y = std::min(std::max(0.f, kps_y), img_height - 1.f);
      box_kps.landmarks.points.push_back(kps);
    }
    box_kps.landmarks.flag = true;
    box_kps.flag = true;

    bbox_kps_collection.push_back(box_kps);

    count += 1; // limit boxes for nms.
    if (count > max_nms)
      break;
  }  
  // printf("[liufeng debug] bbox_kps_collection.size: %d\n",bbox_kps_collection.size());
  if (bbox_kps_collection.size() > nms_pre_)
  {
    std::sort(
        bbox_kps_collection.begin(), bbox_kps_collection.end(),
        [](const BoxfWithLandmarks &a, const BoxfWithLandmarks &b)
        { return a.box.score > b.box.score; }
    ); // sort inplace
    // trunc
    bbox_kps_collection.resize(nms_pre_);
  }

}

void RKScrfdPredictor::generate_bboxes_kps(const SCRFDScaleParams &scale_params,
                             std::vector<BoxfWithLandmarks> &bbox_kps_collection,
                             float** output_tensors,
                             float score_threshold, float img_height,
                             float img_width,float input_height,float input_width)
{
  // score_8,score_16,score_32,bbox_8,bbox_16,bbox_32
  float* score_8 = output_tensors[0];  // e.g [1,12800,1]
  float* score_16 = output_tensors[1]; // e.g [1,3200,1]
  float* score_32 = output_tensors[2]; // e.g [1,800,1]
  float* bbox_8 = output_tensors[3];   // e.g [1,12800,4]
  float* bbox_16 = output_tensors[4];  // e.g [1,3200,4]
  float* bbox_32 = output_tensors[5];  // e.g [1,800,4]

  // generate center points.
  generate_points(input_height, input_width);

  bbox_kps_collection.clear();

  if (use_kps)
  {
    float* kps_8 = output_tensors[6];   // e.g [1,12800,10]
    float* kps_16 = output_tensors[7];  // e.g [1,3200,10]
    float* kps_32 = output_tensors[8];  // e.g [1,800,10]

    // level 8 & 16 & 32 with kps
    generate_bboxes_kps_single_stride(scale_params, score_8, bbox_8, kps_8, 8, score_threshold,
                                            img_height, img_width, bbox_kps_collection);
    // printf("[liufeng debug] level 8: bbox_kps_collection.size:%d\n",bbox_kps_collection.size());
    generate_bboxes_kps_single_stride(scale_params, score_16, bbox_16, kps_16, 16, score_threshold,
                                            img_height, img_width, bbox_kps_collection);
    generate_bboxes_kps_single_stride(scale_params, score_32, bbox_32, kps_32, 32, score_threshold,
                                            img_height, img_width, bbox_kps_collection);
  } // no kps
  else
  {
    // // level 8 & 16 & 32
    // this->generate_bboxes_single_stride(scale_params, score_8, bbox_8, 8, score_threshold,
    //                                     img_height, img_width, bbox_kps_collection);
    // this->generate_bboxes_single_stride(scale_params, score_16, bbox_16, 16, score_threshold,
    //                                     img_height, img_width, bbox_kps_collection);
    // this->generate_bboxes_single_stride(scale_params, score_32, bbox_32, 32, score_threshold,
    //                                     img_height, img_width, bbox_kps_collection);
  }

}                            

float RKScrfdPredictor::iou_of(Boxf tbox,Boxf tb)
{
    
    float inner_x1 = tb.x1 > tbox.x1 ? tb.x1 : tbox.x1;
    float inner_y1 = tb.y1 > tbox.y1 ? tb.y1 : tbox.y1;
    float inner_x2 = tb.x2 < tbox.x2 ? tb.x2 : tbox.x2;
    float inner_y2 = tb.y2 < tbox.y2 ? tb.y2 : tbox.y2;
    float inner_h = inner_y2 - inner_y1 + static_cast<float>(1.0f);
    float inner_w = inner_x2 - inner_x1 + static_cast<float>(1.0f);
    if (inner_h <= static_cast<float>(0.f) || inner_w <= static_cast<float>(0.f))
        return std::numeric_limits<float>::min();
    float inner_area = inner_h * inner_w;
    float tb_area=std::abs<float>(tb.y2 - tb.y1 + 1)*(tb.x2 - tb.x1 + 1);
    float tbox_area=std::abs<float>(tbox.y2 - tbox.y1 + 1)*(tbox.x2 - tbox.x1 + 1);
    return static_cast<float>(inner_area / (tb_area + tbox_area - inner_area));
}

void RKScrfdPredictor::nms_bboxes_kps(std::vector<BoxfWithLandmarks> &input,
                           std::vector<BoxfWithLandmarks> &output,
                           float iou_threshold, unsigned int topk)
{
  if (input.empty()) return;
  std::sort(
      input.begin(), input.end(),
      [](const BoxfWithLandmarks &a, const BoxfWithLandmarks &b)
      { return a.box.score > b.box.score; }
  );
  const unsigned int box_num = input.size();
  std::vector<int> merged(box_num, 0);

  unsigned int count = 0;
  for (unsigned int i = 0; i < box_num; ++i)
  {
    if (merged[i]) continue;
    std::vector<BoxfWithLandmarks> buf;

    buf.push_back(input[i]);
    merged[i] = 1;

    for (unsigned int j = i + 1; j < box_num; ++j)
    {
      if (merged[j]) continue;

      float iou = static_cast<float>(iou_of(input[i].box,input[j].box));

      if (iou > iou_threshold)
      {
        merged[j] = 1;
        buf.push_back(input[j]);
      }

    }
    output.push_back(buf[0]);

    // keep top k
    count += 1;
    if (count >= topk)
      break;
  }
}
//SCRFD人脸检测、关键点检测后处理算法具体函数实现

int RKScrfdPredictor::PostProcess(std::vector<BaseInfo*>& det_infos, 
                int cur_block,
                RKModelManager* mng, 
                int md_idx) {

  return 0;
}

void RKScrfdPredictor::ParseTensor(float* data, int layer, 
  std::vector<cv::Rect>& boxs, std::vector<float>& scores, std::vector<int>& labels,
  std::vector<int> ori_shape, std::vector<int> feteature_shape, ModelConfig* cfg) {
    int num_classes = cfg->label_list.size();
    int stride = cfg->strides[layer];
    int num_pred = num_classes + 4 + 1;
    if (cfg->anchors_table.empty()) {
      spdlog::get("logger")->info("ERROR. RKScrfdPredictor::ParseTensor anchors_table is empty.");
      return;
    }
    std::vector<std::vector<int>> &anchors = cfg->anchors_table[layer];
    float w_scale = ori_shape[1] * 1.0 / cfg->input_shape[1];
    float h_scale = ori_shape[0] * 1.0 / cfg->input_shape[0];
    // spdlog::get("logger")->info("num_classes: {}, stride: {}, num_pred: {}, w_scale: {}, h_scale: {}", num_classes, stride, num_pred, w_scale, h_scale);

    int anchor_num = anchors.size();
    for (int h = 0; h < feteature_shape[0]; h++) {
      for (int w = 0; w < feteature_shape[1]; w++) {
        for (int k = 0; k < anchor_num; k++) {
          float *cur_data = data + k * num_pred;
          // if (h == 0 && w == 0 ) {printf("%f, %f, %f, %f, %f, %f, %f, %f\n", cur_data[0],cur_data[1],cur_data[2],cur_data[3],cur_data[4],cur_data[5],cur_data[6],cur_data[7]);}
          
          int id = std::distance(cur_data + 5, std::max_element(cur_data + 5, cur_data + 5 + num_classes));
          double x1 = 1 / (1 + std::exp(-cur_data[4])) * 1;
          double x2 = 1 / (1 + std::exp(-cur_data[id + 5]));
          double confidence = x1 * x2;
          // if (1 ) {printf("id:%d, 4:%f, id+5:%f, x1:%f, x2:%f, confidence:%f\n", id, cur_data[4], cur_data[id + 5], x1, x2, confidence);}
          if (confidence < cfg->draw_threshold) {continue;}
          // spdlog::get("logger")->info("confidence:{}", confidence);

          double cx = ((1.0 / (1.0 + std::exp(-cur_data[0]))) * 2 - 0.5 + w) * stride;
          double cy = ((1.0 / (1.0 + std::exp(-cur_data[1]))) * 2 - 0.5 + h) * stride;
          double ww = std::pow((1.0 / (1.0 + std::exp(-cur_data[2]))) * 2, 2) * anchors[k][0];
          double hh = std::pow((1.0 / (1.0 + std::exp(-cur_data[3]))) * 2, 2) * anchors[k][1];

          double xmin = (cx - ww / 2.0)* w_scale;
          double ymin = (cy - hh / 2.0)* h_scale;
          double xmax = (cx + ww / 2.0)* w_scale;
          double ymax = (cy + hh / 2.0)* h_scale;

          if (xmax <= 0 || ymax <= 0) { continue; }
          if (xmin > xmax || ymin > ymax) { continue; }
          if (xmin < 0) { xmin = 0;}
          if (ymin < 0) { ymin = 0;}
          if (xmax >= ori_shape[1]) { xmax = ori_shape[1]-1;}
          if (ymax >= ori_shape[0]) { ymax = ori_shape[0]-1;}

          boxs.emplace_back(cv::Rect(xmin, ymin, xmax-xmin, ymax-ymin));
          scores.emplace_back((float)confidence);
          labels.emplace_back(id);
        }
        data = data + num_pred * anchors.size();
      }
    }
}


// int RKScrfdPredictor::get_tensor_hw(hbDNNTensor* tensor, int *height, int *width) {
//   int h_index = 0;
//   int w_index = 0;
//   if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NHWC) {
//     h_index = 1;
//     w_index = 2;
//   }
//   else if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NCHW) {
//     h_index = 2;
//     w_index = 3;
//   }
//   else { return -1; }
//   *height = tensor->properties.validShape.dimensionSize[h_index];
//   *width = tensor->properties.validShape.dimensionSize[w_index];
//   return 0;
// }


}  // namespace hb
